#!/usr/bin/python
'''
Indexes fasta file for project.
'''
import sys
import numpy as np
import fasta_structures
import h5py_driver
import logging

#logging.basicConfig(level=logging.DEBUG, format='[%(levelname)s] (%(threadName)-10s) %(message)s', )
logging.basicConfig(level=logging.INFO, format='[%(levelname)s] (%(threadName)-10s) %(message)s', )
logging.basicConfig(level=logging.ERROR, format='[%(levelname)s] (%(threadName)-10s) %(message)s', )

# callable parameters.
file1 = sys.argv[1]
file2 = sys.argv[2]
file3 = sys.argv[3]
file4 = sys.argv[4]
file5 = sys.argv[5]

# static parameters.
tag1 = "_R3"
tag2 = "_F3"
tag1_sz = len(tag1)
tag2_sz = len(tag2)

name_width = 17
seq_width = 51
init_struct_size = 150
struct_chunk_size = 100
file_chunk_size = 10000

if struct_chunk_size > init_struct_size:
	logging.error("error: chunk size must be < init_size")
	sys.exit(1)

##################### functions.
'''
Reads in a simple fasta file into structure.
'''
def read_fasta(fpath, dt, tmp_arr, hdriver, init, key):

	if init == False:
		h5obj_idx = struct_chunk_size		

	# variable to track position in numpy array.
	struct_idx = 0
	
	# variables to track position in line array.
	line_idx = 0
	lines_sz = 0
	
	# variables used in storing fasta info.
	name = ""
	seq = ""
	
	# open file pointer.
	fin = open(fpath, "rb")
	
	# put fasta information into structure.
	while 1 == 1:
		
		# ensure we have sequence to read.
		if line_idx == lines_sz:
			# read in new data.
			#logging.debug("reading\t%i\t%i" % (chunk_idx, lines_sz))
			lines = fin.readlines(file_chunk_size)
			lines_sz = len(lines)
			#logging.debug("done\t%i" % lines_sz)
			
			# reset pointers.
			line_idx = 0
			
			# check if we are done.
			if lines == []: break
			
		
		# get name.
		if lines[line_idx][0] == ">":
			# ensure not first.
			if name != "":
				# ensure struct is large enough.
				if struct_idx >= tmp_arr.size:
					# See how we handle saving.
					if init == True:
						# save existing data to disk first time.
						hdriver.save_array(tmp_arr)
					else:
						# save old chunk to disk.
						hdriver.update_array(tmp_arr)
											
						# load new chunk from disk.
						logging.debug("retrieve from disk\t%i\t%i" % (h5obj_idx,h5obj_idx + struct_chunk_size))
						tmp_arr = hdriver.retrieve_array(h5obj_idx, h5obj_idx + struct_chunk_size)
						h5obj_idx = h5obj_idx + struct_chunk_size
					
					# move pointer back to 0.
					struct_idx = 0
				
				# copy entry to struct.
				if init == True:
					tmp_arr[struct_idx]['NAME'] = name
					
				else:
					# ensure names match up.
					if name != tmp_arr[struct_idx]['NAME']:
						logging.error("name does not match up: %s\t%s" % (name, tmp_arr[struct_idx]['NAME']))
						sys.exit(1)		
						
					# save seq info.	
					tmp_arr[struct_idx][key] = seq
					
				struct_idx += 1
				
			# reset values.
			name = lines[line_idx][1::].strip()
			seq = ""
			
		else:
			# save seq info.
			seq += lines[line_idx].strip()
			
		# increment line idx.
		line_idx += 1
			
	# save final entry.
	if init == True:
		tmp_arr[struct_idx]['NAME'] = name
		
	else:
		# ensure names match up.
		if name != tmp_arr[struct_idx]['NAME']:
			logging.error("name does not match up: %s\t%s" % (name, tmp_arr[struct_idx]['NAME']))
			sys.exit(1)		
			
		# save seq info.	
		tmp_arr[struct_idx][key] = seq
		
	struct_idx += 1
	
	# shrink array to active size.
	tmp_arr = np.resize(tmp_arr, struct_idx)
	
	# save final array to object.
	hdriver.save_array(tmp_arr)
	fin.close()

'''
Reads in a simple fasta file into structure.
'''
def read_fasta_qual(fpath, qpath, dt, tmp_arr, hdriver):

	# load in the fasta information.
	logging.info("loading fasta info.")
	read_fasta(fpath, dt, tmp_arr, hdriver)

	# load initial chunk from disk and set track variable.
	tmp_arr = hdriver.retrieve_array(0, struct_chunk_size)
	h5obj_idx = struct_chunk_size

	# variable to track position in numpy array.
	struct_idx = 0
	struct_sz = 0
	
	# variables to track position in line array.
	line_idx = 0
	lines_sz = 0
	
	# variables used in storing fasta info.
	name = ""
	seq = ""
	
	# open file pointer.
	fin = open(qpath, "rb")
	
	# put fasta information into structure.
	logging.debug("loading qual info.")
	while 1 == 1:
		
		# ensure we have sequence to read.
		if line_idx == lines_sz:
			# read in new data.
			logging.debug("reading\t%i\t%i" % (line_idx, lines_sz))
			lines = fin.readlines(file_chunk_size)
			lines_sz = len(lines)
			logging.debug("done\t%i" % lines_sz)			
			
			# reset pointers.
			line_idx = 0
			
			# check if we are done.
			if lines == []: break
			
		
		# get name.
		if lines[line_idx][0] == ">":
			
			# ensure not first.
			if name != "":
				# ensure correct chunk is loaded.
				if struct_idx >= tmp_arr.size:
					# save old chunk to disk.
					hdriver.update_array(tmp_arr)
										
					# load new chunk from disk.
					logging.debug("retrieve from disk\t%i\t%i" % (h5obj_idx,h5obj_idx + struct_chunk_size))
					tmp_arr = hdriver.retrieve_array(h5obj_idx, h5obj_idx + struct_chunk_size)
					h5obj_idx = h5obj_idx + struct_chunk_size
					
					# move pointer back to 0.
					struct_idx = 0
				
				# parse out integers.
				tmp = seq.strip().split(" ")
				seq = ""
				for t in tmp:
					seq += chr(int(t))
				
				# ensure names match up.
				if name != tmp_arr[struct_idx]['NAME']:
					logging.error("name does not match up: %s\t%s" % (name, tmp_arr[struct_idx]['NAME']))
					sys.exit(1)
				
				# copy entry to struct.
				#tmp_arr[struct_idx]['NAME'] = name
				tmp_arr[struct_idx]['QUAL1'] = seq
				
				#print name
				struct_idx += 1
				
			# reset values.
			name = lines[line_idx][1::].strip()
			seq = ""
			
		else:
			# save seq info.
			seq += lines[line_idx]
			
		# increment line idx.
		line_idx += 1
					
	# save final entry.
	tmp = seq.strip().split(" ")
	seq = ""
	for t in tmp:
		seq += chr(int(t))
	#tmp_arr[struct_idx]['NAME'] = name
	tmp_arr[struct_idx]['QUAL1'] = seq
	struct_idx += 1
	
	
	# update final array to object.
	hdriver.update_array(tmp_arr)
	fin.close()
	
'''
Reads in paired fasta file into structure.
'''
def read_pair_fasta(fpath1, fpath2, dt, tmp_arr, hdriver):

	# load in the fasta information.
	logging.info("loading fasta1 info.")
	read_fasta(fpath1, dt, tmp_arr, hdriver)

	# load initial chunk from disk and set track variable.
	tmp_arr = hdriver.retrieve_array(0, struct_chunk_size)
	h5obj_idx = struct_chunk_size

	# variable to track position in numpy array.
	struct_idx = 0
	struct_sz = 0
	
	# variables to track position in line array.
	line_idx = 0
	lines_sz = 0
	
	# variables used in storing fasta info.
	name = ""
	seq = ""
	
	# open file pointer.
	fin = open(fpath2, "rb")
	
	# put fasta information into structure.
	logging.info("loading fasta2 info.")
	while 1 == 1:
		
		# ensure we have sequence to read.
		if line_idx == lines_sz:
			# read in new data.
			logging.debug("reading\t%i\t%i" % (line_idx, lines_sz))
			lines = fin.readlines(file_chunk_size)
			lines_sz = len(lines)
			logging.debug("done\t%i" % lines_sz)			
			
			# reset pointers.
			line_idx = 0
			
			# check if we are done.
			if lines == []: break
			
		
		# get name.
		if lines[line_idx][0] == ">":
			
			# ensure not first.
			if name != "":
				# ensure correct chunk is loaded.
				if struct_idx >= tmp_arr.size:
					# save old chunk to disk.
					hdriver.update_array(tmp_arr)
										
					# load new chunk from disk.
					logging.debug("retrieve from disk\t%i\t%i" % (h5obj_idx,h5obj_idx + struct_chunk_size))
					tmp_arr = hdriver.retrieve_array(h5obj_idx, h5obj_idx + struct_chunk_size)
					h5obj_idx = h5obj_idx + struct_chunk_size
					
					# move pointer back to 0.
					struct_idx = 0
				
				# ensure names match up.
				if name[0:-tag2_sz] != tmp_arr[struct_idx]['NAME'][0:-tag1_sz]:
					logging.error("name does not match up: %s\t%s" % (name, tmp_arr[struct_idx]['NAME']))
					sys.exit(1)
				
				# copy entry to struct.
				tmp_arr[struct_idx]['SEQ2'] = seq
				
				#print name
				struct_idx += 1
				
			# reset values.
			name = lines[line_idx][1::].strip()
			seq = ""
			
		else:
			# save seq info.
			seq += lines[line_idx]
			
		# increment line idx.
		line_idx += 1
					
	# save final entry.
	#tmp_arr[struct_idx]['NAME'] = name
	tmp_arr[struct_idx]['SEQ2'] = seq
	struct_idx += 1
	
	
	# update final array to object.
	hdriver.update_array(tmp_arr)
	fin.close()
	
'''
Reads in paired fasta file into structure.
'''
def read_pair_fasta_qual(fpath1, fpath2, qpath1, qpath2, dt, tmp_arr, hdriver):

	# load in the paired fasta information.
	logging.info("loading fasta1 and fasta2 info.")
	read_pair_fasta(fpath1, fpath2, dt, tmp_arr, hdriver)

	

	# load initial chunk from disk and set track variable.
	tmp_arr = hdriver.retrieve_array(0, struct_chunk_size)
	h5obj_idx = struct_chunk_size

	# variable to track position in numpy array.
	struct_idx = 0
	struct_sz = 0
	
	# variables to track position in line array.
	line_idx = 0
	lines_sz = 0
	
	# variables used in storing fasta info.
	name = ""
	seq = ""
	
	# open file pointer.
	fin = open(fpath2, "rb")
	
	# put fasta information into structure.
	logging.info("loading fasta2 info.")
	while 1 == 1:
		
		# ensure we have sequence to read.
		if line_idx == lines_sz:
			# read in new data.
			logging.debug("reading\t%i\t%i" % (line_idx, lines_sz))
			lines = fin.readlines(file_chunk_size)
			lines_sz = len(lines)
			logging.debug("done\t%i" % lines_sz)			
			
			# reset pointers.
			line_idx = 0
			
			# check if we are done.
			if lines == []: break
			
		
		# get name.
		if lines[line_idx][0] == ">":
			
			# ensure not first.
			if name != "":
				# ensure correct chunk is loaded.
				if struct_idx >= tmp_arr.size:
					# save old chunk to disk.
					hdriver.update_array(tmp_arr)
										
					# load new chunk from disk.
					logging.debug("retrieve from disk\t%i\t%i" % (h5obj_idx,h5obj_idx + struct_chunk_size))
					tmp_arr = hdriver.retrieve_array(h5obj_idx, h5obj_idx + struct_chunk_size)
					h5obj_idx = h5obj_idx + struct_chunk_size
					
					# move pointer back to 0.
					struct_idx = 0
				
				# ensure names match up.
				if name[0:-tag2_sz] != tmp_arr[struct_idx]['NAME'][0:-tag1_sz]:
					logging.error("name does not match up: %s\t%s" % (name, tmp_arr[struct_idx]['NAME']))
					sys.exit(1)
				
				# copy entry to struct.
				tmp_arr[struct_idx]['SEQ2'] = seq
				
				#print name
				struct_idx += 1
				
			# reset values.
			name = lines[line_idx][1::].strip()
			seq = ""
			
		else:
			# save seq info.
			seq += lines[line_idx]
			
		# increment line idx.
		line_idx += 1
					
	# save final entry.
	#tmp_arr[struct_idx]['NAME'] = name
	tmp_arr[struct_idx]['SEQ2'] = seq
	struct_idx += 1
	
	
	# update final array to object.
	hdriver.update_array(tmp_arr)
	fin.close()
	
##################### script.


#Test regular fasta
# Initialize temporary array and type.
dt, tmp_arr = fasta_structures.basic_fasta( \
	init_struct_size,\
	nw = name_width,\
	sw = seq_width\
)
	
# open h5py object.
hdriver = h5py_driver.H5_Driver()
hdriver.open_write(file5, "simple_fasta", dt, struct_chunk_size, tmp_arr)

# read fasta file into structure.
read_fasta(file1, dt, tmp_arr, hdriver, True, "SEQ1")


'''
#Test regular fasta qual pair
# Initialize temporary array and type.
dt, tmp_arr = fasta_structures.basic_fasta_qual( \
	init_struct_size,\
	nw = name_width,\
	sw = seq_width\
)
	
# open h5py object.
hdriver = h5py_driver.H5_Driver()
hdriver.open_write(file3, "simple_fasta_qual", dt, struct_chunk_size, tmp_arr)

# read fasta file into structure.
read_fasta_qual(file1, file2, dt, tmp_arr, hdriver)
'''

'''
# Test paired fasta file
# Initialize temporary array and type.
dt, tmp_arr = fasta_structures.paired_fasta( \
	init_struct_size,\
	nw = name_width,\
	sw = seq_width\
)
	
# open h5py object.
hdriver = h5py_driver.H5_Driver()
hdriver.open_write(file3, "paired_fasta", dt, struct_chunk_size, tmp_arr)

# read fasta file into structure.
read_pair_fasta(file1, file2, dt, tmp_arr, hdriver)
'''

'''
# Test paired fasta qual pair
# Initialize temporary array and type.
dt, tmp_arr = fasta_structures.paired_fasta_qual( \
	init_struct_size,\
	nw = name_width,\
	sw = seq_width\
)
	
# open h5py object.
hdriver = h5py_driver.H5_Driver()
hdriver.open_write(file3, "paired_fasta_qual", dt, struct_chunk_size, tmp_arr)

# read fasta file into structure.
read_pair_fasta_qual(file1, file2, file3, file4, dt, tmp_arr, hdriver)
'''
# close h5py object.
hdriver.close()
